#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import rospy
import yaml
from geometry_msgs.msg import PoseStamped

class PointSaver:
    def __init__(self):
        # 获取当前脚本所在目录的路径
        script_dir = os.path.dirname(os.path.abspath(__file__))

        # 构造YAML文件的完整路径 (上一级目录的上一级)
        self.yaml_file = os.path.join(script_dir, "../../../points.yaml")

        # 需要更新的城市名称列表
        self.city_names = [u"深圳", u"上海", u"北京", u"广州", u"吉林"]

        # 初始化当前城市索引
        self.current_index = 0

        # 当前坐标缓存
        self.current_pose = None
        self.received_pose = False  # 标记是否已经接收到有效的位置信息

        # 用于存储 YAML 文件内容的数组
        self.points = []

        # 加载 YAML 文件内容到内存中的 points 数组
        self.points = self.load_yaml(self.yaml_file)
        self.print_points()

        # 初始化 ROS 节点
        rospy.init_node('update_yaml', anonymous=True)

        # 启动订阅器监听当前位置
        rospy.Subscriber('/move_base_simple/goal', PoseStamped, self.pose_callback)

        # 初始提示
        print(u"请移动到 {} 的位置，然后按下 'w' 键保存坐标。".format(self.city_names[self.current_index]).encode('utf-8'))

    def load_yaml(self, file_path):
        print("正在加载YAML文件...")
        with open(file_path, 'r') as file:
            data = yaml.safe_load(file)
        print("YAML文件加载成功。")
        return data

    def save_yaml(self, data, file_path):
        print("正在保存更新后的YAML文件...")
        with open(file_path, 'w') as file:
            yaml.dump(data, file, default_flow_style=False, allow_unicode=True)
        print("YAML文件保存成功。")

    def update_location(self, yaml_data, index, x, y, z, w):
        print("开始更新位置...")
        if index >= len(self.city_names):
            print("索引超出城市列表范围。")
            return

        target_city = self.city_names[index]
        for entry in yaml_data:
            if isinstance(entry, dict) and entry.get("name") == target_city:
                print(u"更新城市: {}".format(target_city).encode('utf-8'))
                entry["x"] = x
                entry["y"] = y
                entry["z"] = z
                entry["w"] = w
                print(u"成功更新城市: {}".format(target_city).encode('utf-8'))
                print("保存的坐标值: x={}, y={}, z={}, w={}".format(x, y, z, w))
                print(u"介绍: {}".format(entry.get("present", "无介绍信息")).encode('utf-8'))
                return

        print(u"未找到城市 {}!".format(target_city).encode('utf-8'))

    def pose_callback(self, msg):
        self.current_pose = msg.pose
        self.received_pose = True  # 设置标记，表示已经接收到位置信息

        print("接收到新的位置信息：x={}, y={}, z={}, w={}\n".format(
            self.current_pose.position.x,
            self.current_pose.position.y,
            self.current_pose.orientation.z,
            self.current_pose.orientation.w
        ))
        print("请输入 'w' 以保存：")

    def print_points(self):
        # 打印加载的 points 内容
        print("加载的 YAML 文件内容:")
        for point in self.points:
            print(point)
            if isinstance(point, dict):
                for key, value in point.items():
                    if isinstance(value, unicode):
                        print(u"{}: {}".format(key, value).encode('utf-8'))
                    else:
                        print("{}: {}".format(key, value))

    def run(self):
        while not rospy.is_shutdown():
            try:
                # 等待用户输入
                user_input = raw_input("请输入 'w' 以保存当前位置：\n").strip()

                # 确保输入为字符串且为 'w'
                if not isinstance(user_input, str) or user_input.lower() != 'w':
                    print("无效输入，请按 'w' 键保存坐标。")
                    continue

                if not self.received_pose:
                    print("尚未接收到有效的位置信息，请先确保机器人移动到指定位置并接收到位置信息。")
                    continue

                if self.current_pose is None or self.current_index >= len(self.city_names):
                    print("当前坐标为空或索引超出范围。")
                    continue

                print("准备保存当前坐标。")
                # 获取当前坐标
                x = self.current_pose.position.x
                y = self.current_pose.position.y
                z = self.current_pose.orientation.z
                w = self.current_pose.orientation.w

                # 更新当前索引对应的城市坐标
                self.update_location(self.points, self.current_index, x, y, z, w)

                # 保存更新后的 YAML 数据
                self.save_yaml(self.points, self.yaml_file)

                # 提示下一个城市或完成信息
                self.current_index += 1
                if self.current_index < len(self.city_names):
                    print(u"请使用rviz导航移动到 {} 的位置，然后按下 'w' 键保存坐标。".format(self.city_names[self.current_index]).encode('utf-8'))
                else:
                    print("所有城市的坐标都已保存完成！")
                    break

            except EOFError as e:
                print("输入读取发生错误: {}".format(e))
            except Exception as e:
                print("发生错误: {}".format(e))

if __name__ == '__main__':
    saver = PointSaver()
    saver.run()

